#!/usr/bin/env python3
import os
os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH", None)  # 确保环境变量被清除

import matplotlib
matplotlib.use('tkagg')  # 确保不和cv2使用qt冲突

from ament_index_python.packages import get_package_share_directory
PATH_SHARE = get_package_share_directory('navigation')
from pathlib import Path
PATH_ROOT = Path(PATH_SHARE).parents[3]
PATH_LIGHTGLUE = PATH_ROOT / "pkgs/LightGlue-ONNX"
import sys
sys.path.append(str(PATH_LIGHTGLUE))
from map_aligner import MapAligner

import rclpy
from rclpy.node import Node
from nav_msgs.msg import OccupancyGrid
from geometry_msgs.msg import Pose, Point, Quaternion
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import numpy as np
import cv2
from threading import Thread
from std_srvs.srv import SetBool

def process_map(grid: OccupancyGrid):
    width, height = grid.info.width, grid.info.height
    map = (np.array(grid.data, dtype=np.double).reshape(height, width) / 100.0 * 255.0).astype(np.uint8)
    map = cv2.cvtColor(map, cv2.COLOR_GRAY2BGR)
    map = np.flipud(map)
    A = np.c_[np.eye(2), np.zeros((2, 1))]
    A[0, 2] = grid.info.origin.position.x / grid.info.resolution
    A[1, 2] = - height - grid.info.origin.position.y / grid.info.resolution
    resolution = grid.info.resolution
    return map, A, resolution

class Map2DAdjuster(Node):
    def __init__(self):
        super().__init__('map_2d_adjuster')
        self.declare_parameter('device', 'cpu')  # 'cuda', 'openvino', 'tensorrt'
        self.declare_parameter('delta_angle', 36.0)
        self.declare_parameter('score_threshold', 80)  # 当匹配点数目大于此值时才认为匹配成功
        device = self.get_parameter('device').get_parameter_value().string_value
        self.map_aligner = MapAligner(device=device)
        self.delta_angle = self.get_parameter('delta_angle').get_parameter_value().double_value
        self.score_threshold = self.get_parameter('score_threshold').get_parameter_value().integer_value

        self.origin_map_subscription = self.create_subscription(OccupancyGrid, 'grid_map', self.origin_map_callback, 1)
        self.origin_map, self.camera_map = None, None
        self.laser_map_subscription = self.create_subscription(OccupancyGrid, 'Laser_2d_map', self.laser_map_callback, 1)
        self.adjust_pose_publisher = self.create_publisher(Pose, 'adjusted_map_pose', 1)
        self.image_pub = self.create_publisher(Image, 'aligned_image', 1)
        self.pause_server = self.create_service(SetBool, 'pause_map_adjuster', self.pause_callback)
        self.bridge = CvBridge()

        self.get_logger().info('Map2DAdjuster has been initialized!')

    def origin_map_callback(self, msg: OccupancyGrid):
        self.origin_map, self.origin_A, self.origin_resolution = process_map(msg)
        self.get_logger().info(f'Received map with {self.origin_map.shape} data points, destroying subscription.')
        self.destroy_subscription(self.origin_map_subscription)
    
    def laser_map_callback(self, msg: OccupancyGrid):
        if self.origin_map is None:
            self.get_logger().warn('Origin map is not set yet, waiting for it.')
            return
        self.camera_map, camera_A, camera_resolution = process_map(msg)
        assert camera_resolution == self.origin_resolution, f"Camera map resolution = {camera_resolution:.2f} must match origin map resolution = {self.origin_resolution:.2f}."
        com_A, score, com_img = self.map_aligner.rotate_compare(
            self.camera_map, self.origin_map, delta_angle=self.delta_angle,
            verbose=False, save=True
        )
        # 发布对齐结果图片
        if com_img is not None:
            img_msg = self.bridge.cv2_to_imgmsg(com_img, encoding="bgr8")
            self.image_pub.publish(img_msg)
        if score < self.score_threshold:
            self.get_logger().warn(f'Not enough matching points: {score} < {self.score_threshold}, skipping pose adjustment.')
            return
        if com_A is None:
            return
        stack = lambda A: np.vstack([A, [0, 0, 1]])
        A = stack(self.origin_A) @ stack(com_A) @ np.linalg.inv(stack(camera_A))  # from camera to origin
        A = np.linalg.inv(A)  # from origin to camera (we need this for ROS transform)
        pose = Pose()
        pose.position = Point(x=A[0, 2] * camera_resolution, y=-A[1, 2] * camera_resolution, z=0.0)
        theta = -np.arctan2(A[1, 0], A[0, 0])
        pose.orientation = Quaternion(
            x=0.0, y=0.0, z=np.sin(theta / 2), w=np.cos(theta / 2)
        )
        self.adjust_pose_publisher.publish(pose)
        self.get_logger().info(f'Publishing adjusted map pose: x = {pose.position.x:.2f}, y = {pose.position.y:.2f}, theta = {theta:.2f} rad')
    
    def pause_callback(self, request: SetBool.Request, response: SetBool.Response):
        if request.data:
            self.get_logger().info('Pausing map adjustment.')
            self.destroy_subscription(self.laser_map_subscription)
        else:
            self.get_logger().info('Starting map adjustment.')
            self.laser_map_subscription = self.create_subscription(OccupancyGrid, 'Laser_2d_map', self.laser_map_callback, 1)
        response.success = True
        response.message = 'Map adjustment paused' if request.data else 'Map adjustment started'
        return response

if __name__ == '__main__':
    rclpy.init()
    rclpy.spin(Map2DAdjuster())
    rclpy.shutdown()
